from contextlib import asynccontextmanager
from typing import AsyncGenerator, TypeVar, Callable, Awaitable

from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker, declarative_base

import logging

MYSQL_USER="zll_rz"
MYSQL_PASSWORD="TRZ58XYytbJafPpj"
MYSQL_HOST="1.14.205.235"
MYSQL_PORT="3306"
MYSQL_DB="zll_rz"
def get_mysql_url() -> str:
    return f"mysql+aiomysql://{MYSQL_USER}:{MYSQL_PASSWORD}@{MYSQL_HOST}:{MYSQL_PORT}/{MYSQL_DB}"

logger = logging.getLogger(__name__)

# 创建 SQLAlchemy 异步引擎
engine = create_async_engine(
    get_mysql_url(),
    echo=True,  # 取消注释以在控制台打印 SQL 语句，用于调试
    pool_recycle=1800,
    pool_size=200,
    max_overflow=500,
    pool_pre_ping=True,
    connect_args={"connect_timeout": 30}

)
# 创建异步会话工厂
AsyncSessionFactory = sessionmaker(
    bind=engine,
    class_=AsyncSession,
    expire_on_commit=False,  # 防止在提交后 ORM 对象过期
)


# 创建数据库引擎和会话工厂(供worker单独使用)
def create_db_engine_and_session_factory():
    engine = create_async_engine(
        get_mysql_url(),
        echo=True,
        pool_recycle=1800,  # 防止超时
        pool_size=200,
        max_overflow=500,
        pool_pre_ping=True,  # 连接前检查火星
    )

    AsyncSessionFactory = sessionmaker(
        bind=engine,
        class_=AsyncSession,
        expire_on_commit=False,
    )
    return engine, AsyncSessionFactory


async def get_db() -> AsyncGenerator[AsyncSession, None]:
    async with AsyncSessionFactory() as session:
        try:
            yield session
        except Exception:
            await session.rollback()
            raise


# 创建一个App上下文管理器，用于在无法传入db参数时执行数据库操作
@asynccontextmanager
async def get_db_context():
    async with AsyncSessionFactory() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()

Base = declarative_base()
